#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import os
import time
from bs4 import BeautifulSoup
import datetime

HOST = 'localhost'
PORT = 27017
DB_NAME = 'pmc'
COLLECTION_NAME = 'metalist'


def GetConfig():
    if os.path.exists("../config.ini"):
        import ConfigParser
        global HOST
        global PORT
        global DB_NAME
        global COLLECTION_NAME
        config = ConfigParser.ConfigParser()
        config.read("../config.ini")
        HOST = config.get("mongodb", "host")
        PORT = int(config.get("mongodb", "port"))
        DB_NAME = config.get("mongodb", "db")
        COLLECTION_NAME = config.get("mongodb", "collection")


def Get_PMCID_Manuscript(article_string):
    soup = BeautifulSoup(article_string, "xml")
    pmcid = None
    manuscriptid = None
    pmcid_node = soup.find("article-id", attrs={"pub-id-type": "pmc"})
    if pmcid_node != None:
        pmcid = pmcid_node.string
    manuscript_node = soup.find("article-id", attrs={"pub-id-type": "manuscript"})
    if manuscript_node != None:
        manuscriptid = manuscript_node.string
    if pmcid != None and pmcid.isdigit():
        return (pmcid + "\t" + str(manuscriptid))
    else:
        return None


# Use beautifulsoup to analyze PMCID lines
def Extract_PMC_IDs_From_XML(filename):
    infile = open(filename, "r")
    outfile = open(filename[:-3] + "ids", "w")
    result = []
    article_string = ""
    is_string_part = False
    for line in infile:
        line = line.rstrip()
        if line.find("<article-meta>") != -1:
            is_string_part = True
        if is_string_part:
            article_string += line
        if line.find("</article-meta>") != -1:
            is_string_part = False
            pmcid_manuscriptid = Get_PMCID_Manuscript(article_string)
            if pmcid_manuscriptid != None:
                result.append(pmcid_manuscriptid)
                outfile.write(pmcid_manuscriptid + "\n")
                outfile.flush()
            article_string = ""
    outfile.close()
    infile.close()
    return result


if __name__ == '__main__':
    print("Loading config file ... "),
    time.sleep(1)
    GetConfig()
    print("complete.")
    time.sleep(1)

    print("Initiating mongodb connection ... "),
    time.sleep(1)
    connection = MongoClient(HOST, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    print("complete.")
    time.sleep(1)

    print("Loading PMC update list ... "),
    time.sleep(1)
    updatelist = {}
    dateString = os.path.abspath(".").split("_")[-1]
    updateListFilename = "pmcid_todo.update." + dateString + ".txt"
    if os.path.getsize(updateListFilename) != 0:
        updatefile = open(updateListFilename, "r")
        for line in updatefile:
            line = line.rstrip()
            splitline = line.split("\t")
            updatelist[int(splitline[0])] = splitline[1]
        updatefile.close()
    print("complete.")
    time.sleep(1)

    print("Loading PMC xml files ... "),
    time.sleep(1)
    filelist = os.listdir(".")
    print("complete.")
    time.sleep(1)

    print("Start anlyzing xml files ... ")
    time.sleep(1)
    file_count = 0
    pmcid_count = 0
    manuscript_total_count = 0
    is_update_file = False
    for filename in filelist:
        if filename.endswith(".xml"):
            filepath = os.path.dirname(os.path.abspath(filename)).split("/")[-1] + "/" + filename
            file_count += 1
            if filename.find("update") != -1:
                is_update_file = True
            else:
                is_update_file = False
            print("Start analyzing " + filename + " ... "),
            ids = Extract_PMC_IDs_From_XML(filename)
            manuscript_count = 0
            pmcid_count += len(ids)
            for id in ids:
                split_id = id.split("\t")
                pmcid = int(split_id[0])
                manuscript_id = split_id[1]
                if manuscript_id != "None":
                    manuscript_count += 1
                    collection.update_one({"pmcid": pmcid}, {"$set": {"xml": filepath, "manuscript_id": manuscript_id}})
                else:
                    collection.update_one({"pmcid": pmcid},
                                          {"$set": {"xml": filepath}, "$unset": {"manuscript_id": True}})
                if is_update_file:
                    if updatelist.has_key(pmcid):
                        updatelist.pop(pmcid)
            manuscript_total_count += manuscript_count
            print("complete. Find " + str(len(ids)) + " PMCIDs, " + str(manuscript_count) + " manuscripts.")
            time.sleep(1)

    print("Reverse no xml update records ... "),
    time.sleep(1)
    if len(updatelist) != 0:
        for key in updatelist:
            if updatelist[key] == "None":
                collection.update_one({"pmcid": pmcid}, {"$unset": {"pmc_date": True}})
            else:
                origin_datetime = datetime.datetime.strptime(updatelist[key], "%Y%m%d%H")
                collection.update_one({"pmcid": pmcid}, {"$set": {"pmc_date": origin_datetime}})
    print(str(len(updatelist)) + " records reversed.")
    time.sleep(1)

    connection.close()
    print("All xml files analyzed. " + str(file_count) + " files, " + str(pmcid_count) + " PMCIDs and " + str(
        manuscript_total_count) + " manuscripts found.")
